import os

os.environ["CUDA_VISIBLE_DEVICES"] = '2,3'
import jsonlines
import traceback
from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download
from modelscope import GenerationConfig
from langchain.prompts import ChatPromptTemplate
import sqlite3
from langchain.utilities import SQLDatabase

# model_dir = '/datasets/fengjiahao/nlp/TongyiFinance/Tongyi-Finance-14B'
# model_dir = '/datasets/fengjiahao/nlp/TongyiFinance/Tongyi-Finance-14B-Chat/'
model_dir = '/datasets/fengjiahao/nlp/qwen/Qwen-7B-Chat/'
# model_dir = '/datasets/fengjiahao/nlp/TongyiFinance/Tongyi-Finance-14B-Chat-Int4/'
# question_json_path = r'/datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/question.json'
question_json_path = r'/datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/submit_result.jsonl'
answer_path = r'/datasets/fengjiahao/nlp/bs_challenge_financial_14b_dataset/submit_result.jsonl'
fund_db_path = r'/public/tmp/fengjiahao/bs_challenge_financial_14b_dataset/dataset/博金杯比赛数据.db'

tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True).eval()
model.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True)  # 可指定不同的生成长度、top_p等相关超参

# conn.configure("busyTimeout", 6000)


content = []
with jsonlines.open(question_json_path, "r") as json_file:
    for obj in json_file.iter(type=dict, skip_invalid=True):
        content.append(obj)


def execute_sql_query(cursor, query):
    result = cursor.execute(query).fetchone()
    return ', '.join(map(str, result)) if result else ''


table_dict = {
    # 'A股': ['A股票日行情表', 'A股股票行业划分表', 'A股股票行情'],
    'A股': ['A股股票行情', '基金股票持仓明细'],
    '基金': ['基金份额持有人结构', '基金债券持仓明细', '基金可转债持仓明细', '基金基本信息', '基金日行情表', '基金股票持仓明细', '基金规模变动表'],
    '港股': ['港股票日行情表', '基金股票持仓明细'],
}
db_info_dict = {}
for theme, table_list in table_dict.items():
    db = SQLDatabase.from_uri(
        "sqlite:///" + fund_db_path,
        include_tables=table_list,
        view_support=True,
        sample_rows_in_table_info=2)
    db_info_dict[theme] = db.table_info

sql_prompt_template = ChatPromptTemplate.from_template(
    "你是一个精通SQL语句的AI。"
    "我会给你一个问题，请按照问题描述，只使用以下表格的表格名和列名写出正确的SQL代码，（数据只是样本，不代表数据库内所有数据，如果SQL中含有日期，需参考样本的格式，最后请按照具体问题的主体和要对应表名和字段名撰写SQL语句）。\n"
    "请直接生成能在以下数据库中执行成功的SQL代码,不要有多余备注和假设,如需表格之间的关联查询请尽量使用LEFT JOIN。\n"
    "请只使用以下数据库的表格名和列名撰写能在SQLite执行的SQL语句(注意字段名与表需打上双引号)：\n\n"
    "{db}\n\n"
    "问题是:请帮我计算，在20210108，中信行业分类划分的一级行业为综合金融行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。"
    "SQL语句：```sql\nSELECT \"股票代码\", ROUND(((\"收盘价(元)\" - \"昨收盘(元)\") / \"昨收盘(元)\") * 100, 2) AS \"涨跌幅\" FROM \"A股股票行情\" WHERE \"一级行业名称\" = '综合金融' AND \"交易日\" = '20210108' ORDER BY \"涨跌幅\" DESC LIMIT 1;\n```\n"
    "问题：{q}\n"
    "SQL语句：\n"
)
sql_prompt_template2 = ChatPromptTemplate.from_template(
    "你是一个精通SQL语句的AI。"
    "我会给你一个问题，请按照问题描述，只使用以下表格的表格名和列名写出正确的SQL代码，（数据只是样本，不代表数据库内所有数据，如果SQL中含有日期，需参考样本的格式，最后请按照具体问题的主体和要对应表名和字段名撰写SQL语句）。\n"
    "请直接生成能在以下数据库中执行成功的SQL代码,不要有多余备注和假设,如需表格之间的关联查询请尽量使用LEFT JOIN。\n"
    "请只使用以下数据库的表格名和列名撰写能在SQLite执行的SQL语句(注意字段名与表需打上双引号)：\n\n"
    "{db}\n\n"
    "问题是:请帮我计算，在20210108，中信行业分类划分的一级行业为综合金融行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。"
    "SQL语句：```sql\nSELECT \"股票代码\", ROUND(((\"收盘价(元)\" - \"昨收盘(元)\") / \"昨收盘(元)\") * 100, 2) AS \"涨跌幅\" FROM \"A股股票行情\" WHERE \"一级行业名称\" = '综合金融' AND \"交易日\" = '20210108' ORDER BY \"涨跌幅\" DESC LIMIT 1;\n```\n"
    "问题：{q}\n"
    "SQL语句：\n"
)
sql_prompt_template3 = ChatPromptTemplate.from_template(
    "你是一个精通SQL语句的AI。"
    "我会给你一个问题，请按照问题描述，只使用以下表格的表格名和列名写出正确的SQL代码，（数据只是样本，不代表数据库内所有数据，如果SQL中含有日期，需参考样本的格式，最后请按照具体问题的主体和要对应表名和字段名撰写SQL语句）。\n"
    "请直接生成能在以下数据库中执行成功的SQL代码,不要有多余备注和假设,如需表格之间的关联查询请尽量使用LEFT JOIN。\n"
    "请只使用以下数据库的表格名和列名撰写能在SQLite执行的SQL语句(注意字段名与表需打上双引号)：\n\n"
    "{db}\n\n"
    "问题是:请帮我计算，在20190129，港股股票里成交金额最高的股票是？成交金额是多少？"
    "SQL语句：```sql\nSELECT t2.\"股票名称\",t1.\"成交金额(元)\" FROM \"港股票日行情表\" AS t1 LEFT JOIN \"基金股票持仓明细\" AS t2 ON t1.\"股票代码\" = t2.\"股票代码\" WHERE t1.\"交易日\"='20190129'  ORDER BY t1.\"成交金额(元)\" DESC LIMIT 1;\n```\n"
    "问题：{q}\n"
    "SQL语句：\n"
)
sql_prompt_template_dict = {
    'A股': sql_prompt_template,
    '基金': sql_prompt_template2,
    '港股': sql_prompt_template3,
}

correct_sql_prompt_template = ChatPromptTemplate.from_template(
    "你回复的SQL语句有语法错误,错误信息如下：\n"
    "{error}\n"
    "记住只使用以下表格的表格名和列名并符合SQLite的语法:\n"
    "{db}\n"
    "根据上面的数据库表格名和列名来修正SQL语句并回答此问题：{q}\n"
    "正确的SQL语句是："
)
answer_prompt_template = ChatPromptTemplate.from_template(
    "你是一个会组织语言的AI。"
    "我会给你一个问题，和相应的答案，请为我完整写出答案句\n"
    "问题：{q}\n"
    "答案：{a}\n"
    "答案句为："
)


# 处理SQL答案
def process_sql_answer(answer):
    try:
        ind = answer.index("SELECT")
        answer = answer[ind:]
        ind = answer.index(';')
        answer = answer[:ind + 1]
    except ValueError:
        pass

    # 去除引号
    if answer[0] in ['"', "'"]:
        answer = answer[1:]
    if answer[-1] == '"':
        answer = answer[:-1]
    if answer[-1] != ';':
        answer += ';'

    # 进行字符串替换
    answer = answer.replace('`', "'")
    answer = answer.replace("\n", " ")

    answer = answer.replace("昨收盘(元)", "昨收盘")
    answer = answer.replace("今开盘(元)", "今开盘")
    answer = answer.replace("最高价(元)", "最高价")
    answer = answer.replace("最低价(元)", "'最低价")
    answer = answer.replace("收盘价(元)", "收盘价")
    answer = answer.replace("成交量(股)", "成交量")
    answer = answer.replace("成交金额(元)", "成交金额")
    answer = answer.replace("所属国家(地区)", "所属国家")

    answer = answer.replace("昨收盘", "昨收盘(元)")
    answer = answer.replace("今开盘", "今开盘(元)")
    answer = answer.replace("最高价", "最高价(元)")
    answer = answer.replace("最低价", "'最低价(元)")
    answer = answer.replace("收盘价", "收盘价(元)")
    answer = answer.replace("成交量", "成交量(股)")
    answer = answer.replace("成交金额", "成交金额(元)")
    answer = answer.replace("所属国家", "所属国家(地区)")

    answer = answer.replace("'昨收盘(元)'", "\"昨收盘(元)\"")
    answer = answer.replace("'今开盘(元)'", "\"今开盘(元)\"")
    answer = answer.replace("'最高价(元)'", "\"最高价(元)\"")
    answer = answer.replace("'最低价(元)'", "\"最低价(元)\"")
    answer = answer.replace("'收盘价(元)'", "\"收盘价(元)\"")
    answer = answer.replace("'成交量(股)'", "\"成交量(股)\"")
    answer = answer.replace("'成交金额(元)'", "\"成交金额(元)\"")
    answer = answer.replace("'所属国家(地区)'", "\"所属国家(地区)\"")

    # answer = answer.replace("昨收盘(元)","'昨收盘(元)'")
    # answer = answer.replace("今开盘(元)", "'今开盘(元)'")
    # answer = answer.replace("最高价(元)", "'最高价(元)'")
    # answer = answer.replace("最低价(元)", "'最低价(元)'")
    # answer = answer.replace("收盘价(元)", "'收盘价(元)'")
    # answer = answer.replace("成交量(股)", "'成交量(股)'")
    # answer = answer.replace("成交金额(元)", "'成交金额(元)'")
    # answer = answer.replace("所属国家(地区)", "'所属国家(地区)'")

    return answer


def ask_llm(ori_question, conn):
    if '公司' not in ori_question or 'A股' in ori_question or '基金' in ori_question or '港股' in ori_question or '股票代码' in ori_question:
        theme = 'A股'
        if '基金' in ori_question:
            theme = '基金'
        elif '港股' in ori_question or '香港' in ori_question:
            theme = '港股'

        # target_table_list = table_dict[theme]
        # table_info = ''
        # for table_name in target_table_list:
        #     table_schema = conn.execute("PRAGMA table_info ( [%s] )"%table_name).fetchall()
        #     table_schema = ', '.join(map(str, table_schema))+'\n'
        #     table_info+=table_schema
        table_info = db_info_dict[theme]
        cur_sql_template = sql_prompt_template_dict[theme]
        prompt = cur_sql_template.format_messages(q=ori_question, db=table_info)
        response, history = model.chat(tokenizer, prompt[0].content, history=None)
        print('!!!response1:', response)
        sql = process_sql_answer(response)
        print('!!!SQL1:', sql)

        cursor = conn.cursor()
        try:
            db_result = execute_sql_query(cursor, sql)
            print('!!!执行SQL成果:', db_result)
        except Exception as e:
            print('!!!执行SQL失败！！')
            # traceback.print_exc()
            print(e)
            print('>>>>>',str(e).split('\n')[-1])
            prompt = correct_sql_prompt_template.format_messages(error=e, db=table_info, q=ori_question)
            response, history = model.chat(tokenizer, prompt[0].content, history=history)
            print('!!!response2:', response)
            sql = process_sql_answer(response)
            print('!!!SQL2:', sql)
            try:
                db_result = execute_sql_query(cursor, sql)
                print('!!!执行SQL成果:', db_result)
            except Exception as e:
                print('!!!执行SQL失败！！')
                # traceback.print_exc()
                print(e)
                print('>>>>>',str(e).split('\n')[-1])
                db_result = '查询数据库失败。'
        try:
            cursor.close()
        except Exception as e:
            traceback.print_exc()
            print(e)
        try:
            prompt = answer_prompt_template.format_messages(q=ori_question, a=db_result)
            response, history = model.chat(tokenizer, prompt[0].content, history=None)
        except Exception as e:
            traceback.print_exc()
            print(e)
    else:
        # q_embedding = embedding_model.encode(ori_question).tolist()
        # results = collection.query(query_texts=ori_question, n_results=1, include=["documents"])
        # ref = ''
        # for result in results["documents"]:
        #     for i in result:
        #         ref += i

        # question = '参考资料如下:%s\n请回答问题:%s\n答案:' % (ref, ori_question)
        # response, history = model.chat(tokenizer, question, history=None)
        response = ''
    print("!!!Q:", question)
    print("!!!A:", response)

    return response


conn = None
for i, cont in enumerate(content):
    if i % 10 == 0:
        # DB
        if conn != None:
            conn.close()
        conn = sqlite3.connect(fund_db_path, timeout=30)
    question = cont['question']
    if 'answer' not in cont or cont['answer'] == '' or '最新的数据' in cont['answer']:
        response = ask_llm(question, conn)
        # print(question)

        cont['answer'] = response
        with jsonlines.open(answer_path, "w") as json_file:
            json_file.write_all(content)
    # 市盈率是最常用来评估股价水平是否合理的指标之一，是指股票价格与每股盈利的比率。...

try:
    conn.close()
except Exception as E:
    print(E)
